{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Using the Segment-Geospatial Python Package with ArcGIS Pro\n",
    "\n",
    "The notebook shows step-by-step instructions for using the Segment Anything Model (SAM) with ArcGIS Pro. Check out the YouTube tutorial [here](https://youtu.be/VvyInoQ6N8Q) and the [Resources for Unlocking the Power of Deep Learning Applications Using ArcGIS](https://community.esri.com/t5/education-blog/resources-for-unlocking-the-power-of-deep-learning/ba-p/1293098). Credit goes to [Esri](https://www.esri.com).\n",
    "\n",
    "[![Alt text](https://img.youtube.com/vi/VvyInoQ6N8Q/0.jpg)](https://youtu.be/VvyInoQ6N8Q)\n",
    "\n",
    "\n",
    "## Installation\n",
    "\n",
    "1. Navigate to the **Start Menu** -> **All apps** -> **ArcGIS** folder, then open the **Python Command Prompt**.\n",
    "2. Create a new conda environment and install dependencies.\n",
    "\n",
    "    `conda create esri::python esri::arcpy conda-forge::segment-geospatial --name geo`\n",
    "\n",
    "3. Activate the new environment in ArcGIS Pro.\n",
    "\n",
    "    `proswap geo`\n",
    "\n",
    "4. Close the Python Command Prompt and open ArcGIS Pro.\n",
    "5. [Download](https://samgeo.gishub.org/examples/arcgis/arcgis.ipynb) this notebook and run it in ArcGIS Pro."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import leafmap\n",
    "from samgeo import SamGeo\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Download sample data\n",
    "\n",
    "In this example, we will use the high-resolution aerial imagery from the USDA National Agricultural Imagery Program ([NAIP](https://naip-usdaonline.hub.arcgis.com/)). You can download NAIP imagery using the [USDA Data Gateway](https://datagateway.nrcs.usda.gov/) or the [USDA NCRS Box Drive](https://nrcs.app.box.com/v/naip). I have downloaded some NAIP imagery and clipped them to a smaller area, which are available [here](https://github.com/opengeos/data/tree/main/naip)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "workspace = os.path.dirname(arcpy.env.workspace)\n",
    "os.chdir(workspace)\n",
    "arcpy.env.overwriteOutput = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "leafmap.download_file(\n",
    "    url=\"https://github.com/opengeos/data/blob/main/naip/buildings.tif\",\n",
    "    quiet=True,\n",
    "    overwrite=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "leafmap.download_file(\n",
    "    url=\"https://github.com/opengeos/data/blob/main/naip/agriculture.tif\",\n",
    "    quiet=True,\n",
    "    overwrite=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "leafmap.download_file(\n",
    "    url=\"https://github.com/opengeos/data/blob/main/naip/water.tif\",\n",
    "    quiet=True,\n",
    "    overwrite=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initialize SAM class\n",
    "\n",
    "Specify the file path to the model checkpoint. If it is not specified, the model will to downloaded to the working directory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sam = SamGeo(\n",
    "    model_type=\"vit_h\",\n",
    "    sam_kwargs=None,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Automatic mask generation\n",
    "\n",
    "Specify the file path to the image we downloaded earlier."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "image = \"agriculture.tif\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can also use your own image. Uncomment and run the following cell to use your own image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# image = '/path/to/your/own/image.tif'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Segment the image and save the results to a GeoTIFF file. Set `unique=True` to assign a unique ID to each object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sam.generate(image, output=\"ag_masks.tif\", foreground=True, unique=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you run into GPU memory errors, uncomment the following code block and run it to empty cuda cache then rerun the code block above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# sam.clear_cuda_cache()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Show the segmentation result as a grayscale image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sam.show_masks(cmap=\"binary_r\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Show the object annotations (objects with random color) on the map."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sam.show_anns(axis=\"off\", alpha=1, output=\"ag_annotations.tif\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Add layers to ArcGIS Pro."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "m = leafmap.arc_active_map()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "m.addDataFromPath(os.path.join(workspace, \"agriculture.tif\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "m.addDataFromPath(os.path.join(workspace, \"ag_annotations.tif\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Convert the object annotations to vector format, such as GeoPackage, Shapefile, or GeoJSON."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "in_raster = os.path.join(workspace, \"ag_masks.tif\")\n",
    "out_shp = os.path.join(workspace, \"ag_masks.shp\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "arcpy.conversion.RasterToPolygon(in_raster, out_shp)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Segment waterbodies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "image = \"water.tif\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sam.generate(image, output=\"water_masks.tif\", foreground=True, unique=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# sam.clear_cuda_cache()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sam.show_masks(cmap=\"binary_r\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sam.show_anns(axis=\"off\", alpha=1, output=\"water_annotations.tif\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "m.addDataFromPath(os.path.join(workspace, \"water.tif\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "m.addDataFromPath(os.path.join(workspace, \"water_annotations.tif\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "in_raster = os.path.join(workspace, \"water_masks.tif\")\n",
    "out_shp = os.path.join(workspace, \"water_masks.shp\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "arcpy.conversion.RasterToPolygon(in_raster, out_shp)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Automatic mask generation options\n",
    "\n",
    "There are several tunable parameters in automatic mask generation that control how densely points are sampled and what the thresholds are for removing low quality or duplicate masks. Additionally, generation can be automatically run on crops of the image to get improved performance on smaller objects, and post-processing can remove stray pixels and holes. Here is an example configuration that samples more masks:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sam_kwargs = {\n",
    "    \"points_per_side\": 32,\n",
    "    \"pred_iou_thresh\": 0.86,\n",
    "    \"stability_score_thresh\": 0.92,\n",
    "    \"crop_n_layers\": 1,\n",
    "    \"crop_n_points_downscale_factor\": 2,\n",
    "    \"min_mask_region_area\": 100,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sam = SamGeo(\n",
    "    model_type=\"vit_h\",\n",
    "    sam_kwargs=sam_kwargs,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sam.generate(\"agriculture.tif\", output=\"ag_masks2.tif\", foreground=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sam.show_masks(cmap=\"binary_r\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sam.show_anns(axis=\"off\", alpha=0.5, output=\"ag_annotations2.tif\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "file_extension": ".py",
   "name": "python",
   "version": "3.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
